#!/bin/bash

# IO监控脚本
# 用于监控系统IO状态并在高负载时发出警告

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# # 检查依赖
# if ! command -v iostat &> /dev/null; then
#     echo -e "${RED}未找到iostat命令${NC}"
#     echo "请安装sysstat包: sudo apt install sysstat"
#     exit 1
# fi

# 默认参数
INTERVAL=5
COUNT=0
THRESHOLD=80
ALERT_MODE=0

# 帮助信息
show_help() {
    echo "IO监控脚本 - 监控系统IO状态并在高负载时发出警告"
    echo
    echo "用法: $0 [选项]"
    echo
    echo "选项:"
    echo "  -i, --interval <秒>    监控间隔，默认5秒"
    echo "  -c, --count <次数>     监控次数，默认无限"
    echo "  -t, --threshold <百分比> IO使用率警告阈值，默认80%"
    echo "  -a, --alert            启用警告模式（高IO时发出声音警告）"
    echo "  -h, --help             显示此帮助信息"
    echo
    echo "监控指标:"
    echo "  - 设备类型 (NVMe/SATA)"
    echo "  - 4K随机读/写速度 (kB/s)"
    echo "  - IO使用率百分比"
    echo
    echo "示例:"
    echo "  $0 -i 2 -t 70 -a       每2秒监控一次，IO使用率超过70%时发出警告"
    echo "  $0 -c 10               监控10次后退出"
}

# 解析参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -i|--interval)
            INTERVAL="$2"
            shift 2
            ;;
        -c|--count)
            COUNT="$2"
            shift 2
            ;;
        -t|--threshold)
            THRESHOLD="$2"
            shift 2
            ;;
        -a|--alert)
            ALERT_MODE=1
            shift
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            echo -e "${RED}未知选项: $1${NC}"
            show_help
            exit 1
            ;;
    esac
done

# 显示监控信息
echo -e "${BLUE}===== IO监控开始 =====${NC}"
echo "监控间隔: ${INTERVAL}秒"
if [ "$COUNT" -gt 0 ]; then
    echo "监控次数: ${COUNT}次"
else
    echo "监控次数: 无限（按Ctrl+C停止）"
fi
echo "警告阈值: ${THRESHOLD}%"
if [ "$ALERT_MODE" -eq 1 ]; then
    echo "警告模式: 已启用"
else
    echo "警告模式: 已禁用"
fi
echo

# 检测存储设备类型
detect_device_type() {
    local device=$1
    if [[ $device == nvme* ]]; then
        echo "NVMe"
    elif [[ $device == sd* ]] || [[ $device == hd* ]]; then
        echo "SATA"
    else
        echo "Unknown"
    fi
}

# 监控函数
monitor_io() {
    local iteration=0
    
    while true; do
        # 获取当前时间
        current_time=$(date "+%Y-%m-%d %H:%M:%S")
        echo -e "${YELLOW}[$current_time] IO状态:${NC}"
        
        # 获取IO统计信息
        io_stats=$(iostat -xd 1 1 | grep -v "^$" | grep -A 1 "^Device" | tail -n +2)
        
        # 显示IO统计信息
        echo "$io_stats" | while read line; do
            if [ -n "$line" ]; then
                device=$(echo "$line" | awk '{print $1}')
                util=$(echo "$line" | awk '{print $NF}')
                
                # 检查是否为数字
                if [[ "$util" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
                    util_int=${util%.*}
                    
                # 获取设备类型
                device_type=$(detect_device_type "$device")
                
                # 获取更详细的IO统计
                detailed_stats=$(iostat -d -k 1 1 $device | tail -n +4 | head -1)
                read_kB=$(echo "$detailed_stats" | awk '{print $4}')
                write_kB=$(echo "$detailed_stats" | awk '{print $5}')
                
                # 显示详细IO信息
                echo -e "${BLUE}设备类型: ${device_type}${NC}"
                echo -e "4K随机读: ${read_kB} kB/s"
                echo -e "4K随机写: ${write_kB} kB/s"
                
                # 根据使用率显示不同颜色
                if [ "$util_int" -ge "$THRESHOLD" ]; then
                    echo -e "${device}: ${RED}${util}%${NC} 使用率"
                    
                    # 如果启用警告模式，发出警告
                    if [ "$ALERT_MODE" -eq 1 ] && [ "$util_int" -ge "$THRESHOLD" ]; then
                        echo -e "\a" # 发出蜂鸣声
                        echo -e "${RED}警告: IO使用率过高!${NC}"
                    fi
                elif [ "$util_int" -ge $(($THRESHOLD * 70 / 100)) ]; then
                    echo -e "${device}: ${YELLOW}${util}%${NC} 使用率"
                else
                    echo -e "${device}: ${GREEN}${util}%${NC} 使用率"
                fi
                else
                    echo -e "${device}: ${util}% 使用率"
                fi
            fi
        done
        
        # 显示等待IO的进程
        io_wait_procs=$(ps aux | awk '$8 ~ /D/ {print $0}')
        if [ -n "$io_wait_procs" ]; then
            echo -e "\n${YELLOW}等待IO的进程:${NC}"
            echo "$io_wait_procs" | head -5
            proc_count=$(echo "$io_wait_procs" | wc -l)
            if [ "$proc_count" -gt 5 ]; then
                echo "... 共 $proc_count 个进程在等待IO"
            fi
        fi
        
        echo
        
        # 增加迭代计数
        ((iteration++))
        
        # 检查是否达到指定次数
        if [ "$COUNT" -gt 0 ] && [ "$iteration" -ge "$COUNT" ]; then
            break
        fi
        
        # 等待下一次监控
        sleep "$INTERVAL"
    done
}

# 开始监控
monitor_io

echo -e "${BLUE}===== IO监控结束 =====${NC}"